
import {
  createContext as createContextOrig,
  useContext as useContextOrig,
  Context as ContextOrig,
  useRef,
  useState,
  useReducer,
  useLayoutEffect,
  createElement,
  MutableRefObject,
  ComponentType,
  ReactNode,
  Provider
} from 'react';
import { unstable_batchedUpdates as batchedUpdates } from 'react-dom';

const CONTEXT_VALUE = Symbol();
const ORIGINAL_PROVIDER = Symbol();


type Version = number;
type Listener<Value> = (
  action: { n: Version, p?: Promise<Value>, v?: Value }
) => void

type ContextValue<Value> = {
  [CONTEXT_VALUE]: {
    /* "v"alue     */ v: MutableRefObject<Value>;
    /* versio"n"   */ n: MutableRefObject<Version>;
    /* "l"isteners */ l: Set<Listener<Value>>;
    /* "u"pdate    */ u: (thunk: () => void, options?: { suspense: boolean }) => void;
  };
};

export interface Context<Value> {
  Provider: ComponentType<{ value: Value; children: ReactNode }>;
  displayName?: string;
}



const createProvider = <Value>(
  ProviderOrig: Provider<ContextValue<Value>>,
) => {
  const ContextProvider = ({ value, children }: { value: Value; children: ReactNode }) => {
    const valueRef = useRef(value);
    const versionRef = useRef(0);
    const [resolve, setResolve] = useState<((v: Value) => void) | null>(null);
    if (resolve) {
      resolve(value);
      setResolve(null);
    }
    const contextValue = useRef<ContextValue<Value>>();
    if (!contextValue.current) {
      const listeners = new Set<Listener<Value>>();
      const update = (thunk: () => void, options?: { suspense: boolean }) => {
        batchedUpdates(() => {
          versionRef.current += 1;
          const action: Parameters<Listener<Value>>[0] = {
            n: versionRef.current,
          };
          // if (options?.suspense) {
          //   action.n *= -1; // this is intentional to make it temporary version
          //   action.p = new Promise<Value>((r) => {
          //     setResolve(() => (v: Value) => {
          //       action.v = v;
          //       delete action.p;
          //       r(v);
          //     });
          //   });
          // }
          listeners.forEach((listener) => listener(action));
          thunk();
        });
      };
      contextValue.current = {
        [CONTEXT_VALUE]: {
          /* "v"alue     */ v: valueRef,
          /* versio"n"   */ n: versionRef,
          /* "l"isteners */ l: listeners,
          /* "u"pdate    */ u: update,
        },
      };
    }
    useLayoutEffect(() => {
      valueRef.current = value;
      versionRef.current += 1;
      (contextValue.current as ContextValue<Value>)[CONTEXT_VALUE].l.forEach((listener) => {
        listener({ n: versionRef.current, v: value });
      });
    }, [value]);
    return createElement(ProviderOrig, { value: contextValue.current }, children);
  };
  return ContextProvider;
};



export function createContext<Value>(defaultValue: Value) {
  const context = createContextOrig<ContextValue<Value>>({
    [CONTEXT_VALUE]: {
      /* "v"alue     */ v: { current: defaultValue },
      /* versio"n"   */ n: { current: -1 },
      /* "l"isteners */ l: new Set(),
      /* "u"pdate    */ u: (f:any) => f(),
    },
  });
  (context as unknown as {
    [ORIGINAL_PROVIDER]: Provider<ContextValue<Value>>;
  })[ORIGINAL_PROVIDER] = context.Provider;
  (context as unknown as Context<Value>).Provider = createProvider(context.Provider);
  delete (context as any).Consumer; // no support for Consumer
  return context as unknown as Context<Value>;
}



export function useContextSelector<Value, Selected>(
  context: Context<Value>,
  selector: (value: Value) => Selected,
) {
  const contextValue = useContextOrig(
    context as unknown as ContextOrig<ContextValue<Value>>,
  )[CONTEXT_VALUE];

  const {
    /* "v"alue     */ v: { current: value },
    /* versio"n"   */ n: { current: version },
    /* "l"isteners */ l: listeners,
  } = contextValue;

  const selected = selector(value);

  const [state, dispatch] = useReducer((
    prev: readonly [Value, Selected],
    action?: Parameters<Listener<Value>>[0],
  ) => {
    if (!action) {
      // case for `dispatch()` below
      return [value, selected] as const;
    }
    if ('p' in action) {
      throw action.p;
    }
    if (action.n === version) {
      if (Object.is(prev[1], selected)) {
        return prev; // bail out
      }
      return [value, selected] as const;
    }
    try {
      if ('v' in action) {
        if (Object.is(prev[0], action.v)) {
          return prev; // do not update
        }
        const nextSelected = selector(action.v);
        if (Object.is(prev[1], nextSelected)) {
          return prev; // do not update
        }
        return [action.v, nextSelected] as const;
      }
    } catch (e) {
      // ignored (stale props or some other reason)
    }
    return [...prev] as const; // schedule update
  }, [value, selected] as const);

  if (!Object.is(state[1], selected)) {
    // schedule re-render
    // this is safe because it's self contained
    dispatch();
  }
  useLayoutEffect(() => {
    listeners.add(dispatch);
    return () => {
      listeners.delete(dispatch);
    };
  }, [listeners]);
  return state[1];
}

